"""
结果可视化脚本（空壳示例）+

功能：
- 读取 config.yaml 中的 plot 配置（如果可用）
- 使用 utils.plot_style 中的 set_chinese_font 和 set_plot_style
- 画一个简单的示例图并保存到 figures/ 目录

此脚本应作为示例或快速验证使用。
"""
from pathlib import Path
import sys
from rich.console import Console

# 确保项目根目录在 sys.path 中，便于导入 utils 等顶级模块
project_root = Path(__file__).resolve().parents[1]
if str(project_root) not in sys.path:
	sys.path.insert(0, str(project_root))

import yaml

# 示例绘图已移除：如需绘图，请在本模块中添加具体绘图函数并调用 utils.plot_style 中的样式函数

from utils.plot_style import force_chinese_font
from src.aux_utils import merge_geo_population

import pandas as pd
import geopandas as gpd
import matplotlib as mpl
import matplotlib.font_manager as fm


def load_plot_config(cfg_path: str = None):
	cfg_path = Path(cfg_path) if cfg_path else Path(__file__).resolve().parents[1] / 'config.yaml'
	if not cfg_path.exists():
		return None
	with cfg_path.open('r', encoding='utf-8') as f:
		cfg = yaml.safe_load(f)
	return cfg.get('plot') if isinstance(cfg, dict) else None


# 示例绘图函数已移除（如需示例，请使用 plot_border_counties_population_heatmap）


def prepare_border_counties_population_data(
	year: int = None,
	sites_cfg_path: str = 'sites.yaml',
	population_parquet: str = None,
	geojson_path: str = None,
	cmap: str = None,
	figsize: tuple = None,
	style_cfg: dict = None,
	scope: str = 'county',  # 'county' or 'city'
):
	"""准备边境县乡镇人口与地理数据，返回用于绘图的数据与元信息。

	返回 dict，包含键：
	- geo: GeoDataFrame（含几何与 '总人口' 列）
	- cmap, vmin, vmax, figsize, dpi, filename, colorbar_bbox, font_family, cfg_plot
	"""
	from pathlib import Path

	console = Console()

	# 中性化：此函数仅准备数据并返回绘图所需的配置信息（包括建议的 font_family），
	# 实际的样式/字体加载应在绘图端完成以便复用不同的绘图场景。

	# 读取 sites.yaml 并准备名称标准化
	border_counties = None
	border_cities = None
	name_suffixes = []
	try:
		if Path(sites_cfg_path).exists():
			with open(sites_cfg_path, 'r', encoding='utf-8') as f:
				cfg = yaml.safe_load(f)
			border_counties = cfg.get('BORDER_COUNTIES') or []
			border_cities = cfg.get('BORDER_CITIES') or []
			if not border_counties:
				border_counties = list((cfg.get('COUNTY_SHORT_NAMES_ALL') or {}).keys())
			name_suffixes = cfg.get('geo', {}).get('name_suffixes_to_strip', []) if isinstance(cfg, dict) else []
	except Exception:
		pass

	def normalize_name(s: str) -> str:
		if s is None:
			return ''
		t = str(s).strip()
		for suf in name_suffixes or []:
			if t.endswith(suf):
				t = t[: -len(suf)]
				break
		return t.replace(' ', '').lower()

	border_counties_norm = [normalize_name(x) for x in (border_counties or [])]
	border_cities_norm = [normalize_name(x) for x in (border_cities or [])]

	# 读取人口数据
	# 优先级：函数参数 > style_cfg(plot) > config.yaml(plot.heatmap) > 内置默认
	# 载入项目级 plot 配置
	cfg_plot = load_plot_config(Path(__file__).resolve().parents[1] / 'config.yaml')
	# heatmap 专用配置
	heat_cfg = None
	if cfg_plot and isinstance(cfg_plot, dict):
		heat_cfg = cfg_plot.get('heatmap') if isinstance(cfg_plot.get('heatmap'), dict) else None
	# 决定使用的值
	population_parquet = population_parquet or (style_cfg.get('population_parquet') if style_cfg and style_cfg.get('population_parquet') else (heat_cfg.get('population_parquet') if heat_cfg else 'data/边境城市七普乡镇.parquet'))
	geojson_path = geojson_path or (style_cfg.get('geojson_path') if style_cfg and style_cfg.get('geojson_path') else (heat_cfg.get('geojson_path') if heat_cfg else 'data/云南乡镇.json'))
	cmap = cmap or (style_cfg.get('cmap') if style_cfg and style_cfg.get('cmap') else (heat_cfg.get('cmap') if heat_cfg else 'Reds'))
	# filename 在此函数中本地定义（prepare 阶段仅返回建议的文件名）
	# 默认使用中文友好名称，基于 scope 与年份
	default_base = '七普_边境县乡镇人口热力图' if scope == 'county' else '七普_边境城市乡镇人口热力图'
	default_year = year or 2020
	default_name = f"{default_base}_{default_year}.png"
	filename = (style_cfg.get('output_filename') if style_cfg and style_cfg.get('output_filename') else (heat_cfg.get('output_filename') if heat_cfg else default_name))
	year = year or (style_cfg.get('default_year') if style_cfg and style_cfg.get('default_year') else (heat_cfg.get('default_year') if heat_cfg else 2020))
	figsize = figsize or (style_cfg.get('figsize') if style_cfg and style_cfg.get('figsize') else (cfg_plot.get('figsize') if cfg_plot and cfg_plot.get('figsize') else (10, 12)))

	try:
		pop_df = pd.read_parquet(population_parquet)
	except Exception as e:
		console.log(f'无法读取人口数据 {population_parquet}: {e}')
		return None

	# 选择年份
	if '年份' in pop_df.columns:
		pop_df_year = pop_df[pop_df['年份'] == year].copy()
	else:
		pop_df_year = pop_df.copy()

	# 标准化 population 中的县名
	if '县' in pop_df_year.columns:
		pop_df_year['county_norm'] = pop_df_year['县'].apply(lambda x: normalize_name(x))

	# 筛选范围（按县或按市）
	if scope == 'city':
		# 尝试从人口表中提取地级市/市字段
		city_col = None
		for cand in ['地名', '地级市', '市', 'city', 'city_name']:
			if cand in pop_df_year.columns:
				city_col = cand
				break
		if city_col:
			pop_df_year['city_norm'] = pop_df_year[city_col].apply(lambda x: normalize_name(x))
		else:
			# 尝试通过县->市映射进行回退（sites.yaml 中的 border_city_non_border_counties）
			county_to_city = {}
			try:
				bmap = cfg.get('border_city_non_border_counties') if isinstance(cfg, dict) else {}
				for city, counties in (bmap or {}).items():
					for c in counties:
						county_to_city[normalize_name(c)] = normalize_name(city)
			except Exception:
				county_to_city = {}
			if 'county_norm' in pop_df_year.columns:
				pop_df_year['city_norm'] = pop_df_year['county_norm'].map(lambda x: county_to_city.get(x))

		if border_cities_norm:
			pop_df_year = pop_df_year[pop_df_year['city_norm'].isin(border_cities_norm)]
		elif 'border_city' in pop_df_year.columns:
			pop_df_year = pop_df_year[pop_df_year['border_city'] == 1]
	else:
		# 筛选边境县
		if border_counties_norm:
			pop_df_year = pop_df_year[pop_df_year['county_norm'].isin(border_counties_norm)]
		elif 'border_county' in pop_df_year.columns:
			pop_df_year = pop_df_year[pop_df_year['border_county'] == 1]

	if pop_df_year.empty:
		console.log('未找到匹配的边境县人口数据，绘图终止。')
		return None

	# 准备合并字段
	merge_on_code = None
	if '乡镇代码' in pop_df_year.columns:
		pop_df_year['town_code'] = pop_df_year['乡镇代码'].astype(str)
		merge_on_code = 'town_code'
	elif 'tid' in pop_df_year.columns or 'town_code' in pop_df_year.columns:
		col = 'tid' if 'tid' in pop_df_year.columns else 'town_code'
		pop_df_year['town_code'] = pop_df_year[col].astype(str)
		merge_on_code = 'town_code'

	try:
		gj = gpd.read_file(geojson_path)
	except Exception as e:
		console.log(f'读取 GeoJSON 失败：{e}')
		return None

	geo = gj.copy()
	# 乡镇代码或名称
	if '乡镇代码' in geo.columns:
		geo['town_code'] = geo['乡镇代码'].astype(str)
	elif 'tid' in geo.columns:
		geo['town_code'] = geo['tid'].astype(str)
	else:
		if '乡' in geo.columns:
			geo['town_name'] = geo['乡'].astype(str)
		elif '乡镇' in geo.columns:
			geo['town_name'] = geo['乡镇'].astype(str)
		elif '名称' in geo.columns:
			geo['town_name'] = geo['名称'].astype(str)

	# 创建 county_norm 或 city_norm
	for cand in ['县', 'county', 'COUNTY', '县级']:
		if cand in geo.columns:
			geo['county_norm'] = geo[cand].apply(lambda x: normalize_name(x))
			break
	# city 列检测（多候选项）
	for cand in ['地名', '地级市', '市', 'city', 'city_name']:
		if cand in geo.columns:
			geo['city_norm'] = geo[cand].apply(lambda x: normalize_name(x))
			break

	if scope == 'city':
		# 如果没有 city_norm，则尝试通过 county->city 映射回退
		if 'city_norm' not in geo.columns:
			try:
				bmap = cfg.get('border_city_non_border_counties') if isinstance(cfg, dict) else {}
				county_to_city = {normalize_name(c): normalize_name(city) for city, lst in (bmap or {}).items() for c in lst}
				if 'county_norm' in geo.columns:
					geo['city_norm'] = geo['county_norm'].map(lambda x: county_to_city.get(x))
			except Exception:
				pass
		if border_cities_norm and 'city_norm' in geo.columns:
			geo = geo[geo['city_norm'].isin(border_cities_norm)].copy()
	else:
		if border_counties_norm and 'county_norm' in geo.columns:
			geo = geo[geo['county_norm'].isin(border_counties_norm)].copy()

	merged = merge_geo_population(geo, pop_df_year, merge_on_code, name_suffixes, normalize_name)

	cfg_plot = load_plot_config(Path(__file__).resolve().parents[1] / 'config.yaml')
	plot_dpi = style_cfg.get('dpi') if style_cfg and style_cfg.get('dpi') else (cfg_plot.get('dpi') if cfg_plot and cfg_plot.get('dpi') else 300)
	plot_figsize = style_cfg.get('figsize') if style_cfg and style_cfg.get('figsize') else (cfg_plot.get('figsize') if cfg_plot and cfg_plot.get('figsize') else figsize)

	# 计算绘图所需的参数（但不在此处绘图或保存）
	vmin = merged['总人口'].min()
	vmax = merged['总人口'].max()
	cb_bbox = style_cfg.get('colorbar_bbox') if style_cfg and style_cfg.get('colorbar_bbox') else (heat_cfg.get('colorbar_bbox') if heat_cfg and heat_cfg.get('colorbar_bbox') else [0.40, 0.74, 0.44, 0.035])

	# 建议字体（来自 style_cfg 或 cfg_plot），但不在此处加载字体
	suggested_font = None
	if style_cfg and style_cfg.get('font_family'):
		suggested_font = style_cfg.get('font_family')
	elif cfg_plot and isinstance(cfg_plot, dict) and cfg_plot.get('font_family'):
		suggested_font = cfg_plot.get('font_family')

	# 检查字体是否可用（不修改全局 rc）
	font_available = False
	available_fonts = []
	try:
		seen = set()
		for f in fm.fontManager.ttflist:
			if f.name not in seen:
				available_fonts.append(f.name)
				seen.add(f.name)
			if suggested_font and f.name == suggested_font:
				font_available = True
			if len(available_fonts) >= 10:
				break
	except Exception:
		# 如果无法访问 fontManager，则保持空列表/不可用
		available_fonts = []

	return {
		'geo': merged,
		'cmap': cmap,
		'vmin': vmin,
		'vmax': vmax,
		'figsize': plot_figsize,
		'dpi': plot_dpi,
		'filename': filename,
		'colorbar_bbox': cb_bbox,
		'font_family': suggested_font,
		'font_available': font_available,
		'available_fonts': available_fonts,
		'cfg_plot': cfg_plot,
	}


def main():
	# 尝试从项目根读取 config.yaml
	repo_root = Path(__file__).resolve().parents[1]
	cfg_path = repo_root / 'config.yaml'
	style_cfg = load_plot_config(cfg_path) if cfg_path.exists() else None

	figures_dir = repo_root / (style_cfg.get('figures_dir') if style_cfg and style_cfg.get('figures_dir') else 'figures')
	# 默认仅输出占位示例。按需取消注释下面的行以生成边境县人口热力图。
	# make_sample_plot(out_dir=figures_dir, style_cfg=style_cfg)

	# 当脚本通过命令行运行并带有 --heatmap 参数时，生成热力图
	import argparse
	parser = argparse.ArgumentParser(description='plot utilities (data preparation only)')
	parser.add_argument('--info', action='store_true', help='打印提示：该模块仅负责数据准备，绘图请使用 try.py')
	args = parser.parse_args()
	if args.info:
		print('此模块已重构为仅准备数据。请在 try.py 中调用 prepare_border_counties_population_data 并执行绘图/保存。')


if __name__ == '__main__':
	main()

